#include "stdafx.h"
#include "cmdline.h"
#ifndef _WIN32
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <dirent.h>
#endif

#ifndef _WIN32
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <dirent.h>
#endif

using namespace xm;

struct THIS_MODULE_EXE
{
	tcs::xstr_t ThisDir;
	tcs::xstr_t	ThisPath;
	
	THIS_MODULE_EXE()
	{
		std::string tmp; tmp.resize(2048);
#ifdef _WIN32
		DWORD len = ::GetModuleFileNameA( GetModuleHandleA(0), &tmp[0], tmp.size() );
#else
		DWORD len = ::readlink("/proc/self/exe", &tmp[0], tmp.size() - 1);
		if(len <= 0 || len == tmp.size() - 1) {
		  // fail
		  return;
		}
		tmp[len] = 0;
#endif
		ThisPath.set( 0, tmp.c_str() );
#ifdef _WIN32
		char * p = tcsrchr( tmp.c_str(), '\\' );
#else
		char * p = tcsrchr( tmp.c_str(), '/' );
#endif
		if( p ) *p = '\x00';
		ThisDir.set( 0, tmp.c_str() );
	}
};

THIS_MODULE_EXE _ThisExe;

struct logger
{
	bool				_stdout;
	std::vector<FILE*>	_logfiles;
	logger()
	{
		makeSilence(false);
	}
	bool makeSilence( bool f )
	{
		_stdout = (f == false);
		return f;
	}
	bool addLogFile( const char * fname )
	{
		if(!fname) return false;
		FILE* f = fopen( fname, "w" );
		if( !f ) return false;
		_logfiles.push_back(f);
		return true;
	}
	void _vprintf( const char * format, va_list ap )
	{
		if( _stdout ) vfprintf( stdout, format, ap );
		size_t c = _logfiles.size();
		for( size_t i = 0; i < c; ++ i )
		{
			vfprintf( _logfiles[i], format, ap );
			fflush( _logfiles[i] );
		}
	}
	void stdprintf( const char * format, ... )
	{
		va_list al;
		va_start( al, format);
		vfprintf( stdout, format, al );
	}
	void closeAll()
	{
		while( _logfiles.size() )
		{
			fclose( _logfiles.back() );
			_logfiles.pop_back();
		}
	}
};

logger Logger;

bool logprintf( const char * fmt, ... )
{
	if( !fmt ) return false;
	va_list al;
	va_start( al, fmt);
	Logger._vprintf( fmt, al );
	return true;
}

#define conprintf	printf

struct XManEngine
{
	iptr<IKXoTable>			_got;
	iptr<IKDetectEngine>	_engine;
	size_t					_count;

	XManEngine() : _count(0)
	{}

	~XManEngine()
	{
		if( _got ) _got->Clear();
	}

	HRESULT Init( IKArguments * argsp )
	{
		RFAILED( _got.instantiation( PROGID_XoTable ) );
		RFAILED( _engine.instantiation( PROGID_XdoChain, 0, _got ) );
		if( argsp ) _got->SetUnknown( "arguments", argsp );
		return S_OK;
	}
	VOID ScanOneFile( LPCSTR filename )
	{
		std::string vir_names;
		_ScanOneFile( filename, vir_names ) ;
		if( vir_names.empty() ) 
		{
			//return;
			SetConsoleTextAttribute( GetStdHandle(STD_OUTPUT_HANDLE), FOREGROUND_GREEN|FOREGROUND_INTENSITY );
		}
		else
		{
			SetConsoleTextAttribute( GetStdHandle(STD_OUTPUT_HANDLE), FOREGROUND_RED|FOREGROUND_INTENSITY );
			_count ++;
		}
		printf( "%s    [%s]\n", filename, vir_names.empty() ? "Clean" : vir_names.c_str() );
	}

	HRESULT _ScanOneFile( LPCSTR filename, std::string& virname )
	{
		iptr<IKNativeFile> oHandleFile( PROGID_WinNativeFile );
		iptr<IKVirtalFile> oCacheFile( PROGID_VirtualFile );
		iptr<IKFormats> oFormats( PROGID_KFormats );
		iptr<IKDetectSession> oSession( PROGID_DetectSession );
		if( !oHandleFile || !oCacheFile || !oFormats || !oSession ) 
			return E_FAIL;
		RFAILED( oHandleFile->Open( filename, F_READ|F_OPEN ) );
		RFAILED( oCacheFile->Open( IKVirtalFile::MODE_CACHEOFBDO, F_READ, oHandleFile, 0 ) );
		iptr<IKFileStream> oFile( oCacheFile );
		RFAILED( XmFormatDiscern( oFile, oFormats ) );
		RFAILED( oSession->Open( oCacheFile, DOSCAN_LOGDR ) );
		for( size_t i = 0; i < oFormats->Count(); ++ i )
		{
			oSession->SelectFormat( oFormats->GetAt(i) );
			DETCC cc = _engine->Invoke( oSession );
		}
		size_t numOfResults = oSession->EnumResult( 0, 0 );
		if( !numOfResults ) return S_OK;

		std::vector<DetectResult> oResults;
		oResults.resize( numOfResults );

		numOfResults = oSession->EnumResult( oResults.data(), numOfResults );

		for( size_t i = 0; i < numOfResults; ++ i )
		{
			DetectResult & rResult = oResults[i];
			tcs::xprintf( virname, " %d,%s%s", rResult.SigId, rResult.Level ? "" : "Heuristic:", rResult.Name );
		}
		return S_OK;
	}

	HRESULT ScanDir( LPCSTR wDir )
	{
#ifdef WIN32
		WIN32_FIND_DATAA oFindData = {};
		std::string aFindStr = wDir;
		aFindStr.append( "\\*" );
		HANDLE hFind = FindFirstFileA( aFindStr.c_str(), &oFindData );
		do {

			if( oFindData.dwFileAttributes & FILE_ATTRIBUTE_REPARSE_POINT ) 
				continue;

			if( oFindData.cFileName[0] == '.' )
				continue;

			std::string wName = wDir;
			wName.push_back( '\\' );
			wName.append( oFindData.cFileName );

			if( oFindData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY )
				ScanDir( wName.c_str() );
			else
				ScanOneFile( wName.c_str() );

		} while( FindNextFileA( hFind, &oFindData) );

		FindClose( hFind );
		
		return S_OK;
#else
		DIR *dirp;  
		struct dirent *direntp;  
		if((dirp = opendir(wDir)) == NULL)
		  return E_FAIL;
		std::string filePath;
		while((direntp = readdir(dirp)) != NULL) {
		  if(strcmp(direntp->d_name, ".") == 0 || strcmp(direntp->d_name, "..") == 0)
		    continue;
		  filePath = wDir;
		  filePath.append("/");
		  filePath.append(direntp->d_name);
		  ScanMix(filePath.c_str());
		}
		closedir(dirp);
		return S_OK;
#endif	
	}

	HRESULT ScanMix( LPCSTR wDir )
	{
#ifdef WIN32
		DWORD dwFileAttr = GetFileAttributesA( wDir );
		if( dwFileAttr == INVALID_FILE_ATTRIBUTES ) return E_FAIL;
		if( dwFileAttr & FILE_ATTRIBUTE_REPARSE_POINT ) return E_FAIL;
		if( dwFileAttr & FILE_ATTRIBUTE_DIRECTORY )
			ScanDir( wDir );
		 ScanOneFile( wDir );
		 return S_OK;
#else
		struct stat statInfo;
		if(stat(wDir, &statInfo) != 0) return E_FAIL;
		if(S_ISREG(statInfo.st_mode))
		  return ScanOneFile(wDir);
		else if(S_ISDIR(statInfo.st_mode))
		  return ScanDir(wDir);
		else
		  return E_FAIL;
#endif
	}
};

int main(int argc, char* argv[])
{
	setlocale(LC_ALL, "zh_CN.GB2312");
	//////////////////////////////////////////////////////////////////////////
	iptr<xm::IKArguments> args( PROGID_Arguments );
	if( !args ) return E_FAIL;
	for( int i = 1; i < argc; ++ i )
	{
		if( argv[i][0] == '-' )
			args->SetArgvByString( argv[i] + 1 );
	}

	XManEngine engine;
	RFAILED( engine.Init(args) );

	for( int i = 1; i < argc; ++ i )
	{
		if( argv[i][0] != '-' )
			engine.ScanMix( argv[i] );
	}

	printf( "malware: %d\n", engine._count );

	//clock_t endTime = clock();
	//size_t _mseconds = endTime%1000;
	//size_t _seconds = endTime/1000;
	//size_t _minutes = _seconds/60; _seconds = _seconds%60;
	//logprintf( "TotalCast: %d:%d:%d(m:s:ms)\n", _minutes, _seconds, _mseconds );
	return 0;
};
